import vtkmodules.all as vtk
import numpy as np


class BdfLoader:
    url = ""
    origin = None
    # 记录每个索引位置对应的pointId key: pointId, value: index
    points_id_index_map = None
    # 记录了每个eid的索引位置 key: eid, value: index
    elements_id_index_map = None
    # 记录了索引位置处cell对应的point的索引 key: index, value: [pointIndex1, pointIndex2, pointIndex3 ...]
    cells_id_index_map = None

    # 区分开是为了更好的展示性能
    # 不需要一次性全部展示出来
    # 同时方便3d的进行表面提取

    # 二维多边形单元
    poly_data = None
    points = None
    cells = None

    # 二维线性单元
    line_data = None
    line_index_arr = None
    line_origin_data = None
    line_id_arr = None
    line_poly_data = None
    line_points = None
    line_cells = None

    # 三维单元
    d3_poly_data = None
    d3_points = None
    d3_cells = None

    # 类型map
    type_map = None

    # 物理区域
    blocks_map = None

    has_line = False
    has_d3 = False
    has_d2 = False

    def __init__(self, url):
        self.url = url
        self.poly_data = vtk.vtkPolyData()
        self.points = vtk.vtkPoints()
        self.cells = vtk.vtkCellArray()
        self.points_id_index_map = {}
        self.elements_id_index_map = {}
        self.cells_id_index_map = {}
        self.line_data = []
        self.line_index_arr = []
        self.line_origin_data = []
        self.line_id_arr = []
        self.line_poly_data = vtk.vtkPolyData()
        self.line_points = vtk.vtkPoints()
        self.line_cells = vtk.vtkCellArray()
        self.d3_poly_data = vtk.vtkPolyData()
        self.d3_points = vtk.vtkPoints()
        self.d3_cells = vtk.vtkCellArray()
        self.type_map = {}
        self.blocks_map = {}

    def parse(self, data_obj):
        self.origin = data_obj
        self.create_id_index_map(data_obj)
        self.create_point_data(data_obj)
        self.create_cell_data(data_obj)
        self.poly_data.SetPoints(self.points)
        self.poly_data.SetPolys(self.cells)
        self.d3_poly_data.SetPoints(self.points)
        self.d3_poly_data.SetPolys(self.d3_cells)
        if len(self.line_data) > 0:
            self.has_line = True
            self._create_lines_poly_data(self.line_data)

    def create_id_index_map(self, data_obj):
        ids = data_obj.get("nodeIds")
        i = 0
        for pid in ids:
            self.points_id_index_map[pid] = i
            i += 1

    def create_point_data(self, data_obj):
        nodes = data_obj.get("nodes")
        for i in range(0, int(len(nodes) / 3)):
            self.points.InsertNextPoint(
                [nodes[3 * i], nodes[3 * i + 1], nodes[3 * i + 2]]
            )

    def create_cell_data(self, data_obj):
        tmp_cells = data_obj.get("elements")
        # 因为有单元是多个三角组成的 所以需要偏移处理
        offset = 0
        d3_offset = 0
        line_offset = 0
        for i in range(0, len(tmp_cells)):
            ids_arr = []
            for pid in tmp_cells[i]["nodeIds"]:
                ids_arr.append(self.points_id_index_map[pid])
            self.cells_id_index_map[i] = ids_arr
            point_count = len(ids_arr)
            # 区分线性单元
            if point_count == 2:
                self.line_data.append(data_obj["nodes"][3 * ids_arr[0]])
                self.line_data.append(data_obj["nodes"][3 * ids_arr[0] + 1])
                self.line_data.append(data_obj["nodes"][3 * ids_arr[0] + 2])
                self.line_data.append(data_obj["nodes"][3 * ids_arr[1]])
                self.line_data.append(data_obj["nodes"][3 * ids_arr[1] + 1])
                self.line_data.append(data_obj["nodes"][3 * ids_arr[1] + 2])
                self.line_index_arr.extend(ids_arr)
                self.line_id_arr.append(tmp_cells[i]["eid"])
                self.line_origin_data.append(tmp_cells[i])
                self._inject_cell_to_block_map(
                    tmp_cells[i]["physicalTag"], "line", line_offset
                )
                line_offset += 1
            elif point_count == 3 or point_count == 4:
                self.elements_id_index_map[tmp_cells[i]["eid"]] = offset
                self.cells.InsertNextCell(point_count, ids_arr)
                self._inject_cell_to_map(tmp_cells[i]["type"], offset)
                self._inject_cell_to_block_map(
                    tmp_cells[i]["physicalTag"], "d2", offset
                )
                offset += 1
                self.has_d2 = True
            else:
                self.has_d3 = True
                match tmp_cells[i]["type"]:
                    case "CTETRA":
                        self._parse_ctetra(tmp_cells[i], ids_arr)
                        self._inject_cell_to_map("CTETRA", d3_offset)
                        self._inject_cell_to_map("CTETRA", d3_offset + 1)
                        self._inject_cell_to_map("CTETRA", d3_offset + 2)
                        self._inject_cell_to_map("CTETRA", d3_offset + 3)
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset
                        )
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset + 1
                        )
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset + 2
                        )
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset + 3
                        )
                        d3_offset += 4
                    case "CHEXA":
                        self._parse_chexa(tmp_cells[i], ids_arr)
                        self._inject_cell_to_map("CHEXA", d3_offset)
                        self._inject_cell_to_map("CHEXA", d3_offset + 1)
                        self._inject_cell_to_map("CHEXA", d3_offset + 2)
                        self._inject_cell_to_map("CHEXA", d3_offset + 3)
                        self._inject_cell_to_map("CHEXA", d3_offset + 4)
                        self._inject_cell_to_map("CHEXA", d3_offset + 5)
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset
                        )
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset + 1
                        )
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset + 2
                        )
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset + 3
                        )
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset + 4
                        )
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset + 5
                        )
                        d3_offset += 6
                    case "CPENTA":
                        self._parse_cpenta(tmp_cells[i], ids_arr)
                        self._inject_cell_to_map("CHEXA", d3_offset)
                        self._inject_cell_to_map("CHEXA", d3_offset + 1)
                        self._inject_cell_to_map("CHEXA", d3_offset + 2)
                        self._inject_cell_to_map("CHEXA", d3_offset + 3)
                        self._inject_cell_to_map("CHEXA", d3_offset + 4)
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset
                        )
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset + 1
                        )
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset + 2
                        )
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset + 3
                        )
                        self._inject_cell_to_block_map(
                            tmp_cells[i]["physicalTag"], "d3", d3_offset + 4
                        )
                        d3_offset += 5

    def cell_index_to_cell_id(self, cell_index) -> int:
        for key in self.elements_id_index_map:
            if cell_index == self.elements_id_index_map[key]:
                return key
        return -1

    def point_index_to_point_coordinate(self, point_index: int) -> [int]:
        x = self.origin.get("nodes")[point_index * 3]
        y = self.origin.get("nodes")[point_index * 3 + 1]
        z = self.origin.get("nodes")[point_index * 3 + 2]
        return [x, y, z]

    def _parse_multi_face(self, cell_info, ids_arr):
        faces = cell_info["faces"]
        for face in faces:
            face_ids_arr = []
            for id_index in face:
                face_ids_arr.append(ids_arr[id_index])
            self.d3_cells.InsertNextCell(len(face_ids_arr), face_ids_arr)

    def _parse_ctetra(self, cell_info, ids_arr):
        self._parse_multi_face(cell_info, ids_arr)

    def _parse_chexa(self, cell_info, ids_arr):
        self._parse_multi_face(cell_info, ids_arr)

    def _parse_cpenta(self, cell_info, ids_arr):
        self._parse_multi_face(cell_info, ids_arr)

    # 将cell 索引注入到map表中
    def _inject_cell_to_map(self, cell_type: str, index: int):
        type_index_arr = self.type_map.get(cell_type, None)
        if type_index_arr is None:
            type_index_arr = []
            self.type_map[cell_type] = type_index_arr
        type_index_arr.append(index)

    def _inject_cell_to_block_map(self, physicalTag, dim_type, index: int):
        block_index_group = self.blocks_map.get(physicalTag, None)
        if block_index_group is None:
            block_index_group = {
                "line": [],
                "d2": [],
                "d3": [],
            }
            self.blocks_map[physicalTag] = block_index_group
        if dim_type == "line":
            block_index_group["line"].append(index)
        elif dim_type == "d2":
            block_index_group["d2"].append(index)
        elif dim_type == "d3":
            block_index_group["d3"].append(index)

    def _create_lines_cells(self, count):
        for i in range(count):
            self.line_cells.InsertNextCell(2, [i * 2, i * 2 + 1])
            self._inject_cell_to_map("CBAR", i)

    def _create_lines_poly_data(self, points):
        tmp = []
        tmp.extend(points)
        for i in range(0, int(len(tmp) / 3)):
            self.line_points.InsertNextPoint(
                [tmp[3 * i], tmp[3 * i + 1], tmp[3 * i + 2]]
            )
        self._create_lines_cells(int(len(tmp) / 6))
        self.line_poly_data.SetPoints(self.line_points)
        self.line_poly_data.SetLines(self.line_cells)
